{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "numpy.dtype size changed, may indicate binary incompatibility. Expected 96 from C header, got 88 from PyObject",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 6\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mos\u001b[39;00m \n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mnumpy\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m----> 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01msklearn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodel_selection\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m KFold\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mdatasets\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Dataset\n",
      "File \u001b[1;32mc:\\Users\\15144\\miniconda3\\envs\\d2l\\lib\\site-packages\\pandas\\__init__.py:29\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnumpy\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[0;32m     23\u001b[0m     np_version_under1p17 \u001b[38;5;28;01mas\u001b[39;00m _np_version_under1p17,\n\u001b[0;32m     24\u001b[0m     np_version_under1p18 \u001b[38;5;28;01mas\u001b[39;00m _np_version_under1p18,\n\u001b[0;32m     25\u001b[0m     is_numpy_dev \u001b[38;5;28;01mas\u001b[39;00m _is_numpy_dev,\n\u001b[0;32m     26\u001b[0m )\n\u001b[0;32m     28\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m---> 29\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_libs\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m hashtable \u001b[38;5;28;01mas\u001b[39;00m _hashtable, lib \u001b[38;5;28;01mas\u001b[39;00m _lib, tslib \u001b[38;5;28;01mas\u001b[39;00m _tslib\n\u001b[0;32m     30\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:  \u001b[38;5;66;03m# pragma: no cover\u001b[39;00m\n\u001b[0;32m     31\u001b[0m     \u001b[38;5;66;03m# hack but overkill to use re\u001b[39;00m\n\u001b[0;32m     32\u001b[0m     module \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mstr\u001b[39m(e)\u001b[38;5;241m.\u001b[39mreplace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcannot import name \u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\15144\\miniconda3\\envs\\d2l\\lib\\site-packages\\pandas\\_libs\\__init__.py:13\u001b[0m\n\u001b[0;32m      1\u001b[0m __all__ \u001b[38;5;241m=\u001b[39m [\n\u001b[0;32m      2\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNaT\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m      3\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNaTType\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m      9\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInterval\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m     10\u001b[0m ]\n\u001b[1;32m---> 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_libs\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minterval\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Interval\n\u001b[0;32m     14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_libs\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtslibs\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[0;32m     15\u001b[0m     NaT,\n\u001b[0;32m     16\u001b[0m     NaTType,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     21\u001b[0m     iNaT,\n\u001b[0;32m     22\u001b[0m )\n",
      "File \u001b[1;32mpandas\\_libs\\interval.pyx:1\u001b[0m, in \u001b[0;36minit pandas._libs.interval\u001b[1;34m()\u001b[0m\n",
      "\u001b[1;31mValueError\u001b[0m: numpy.dtype size changed, may indicate binary incompatibility. Expected 96 from C header, got 88 from PyObject"
     ]
    }
   ],
   "source": [
    "from typing import Dict, Optional\n",
    "import time\n",
    "import os \n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import KFold\n",
    "from datasets import Dataset\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "from datasets import Dataset, load_dataset\n",
    "from transformers import  TrainingArguments,DataCollatorForLanguageModeling\n",
    "from trl import DPOTrainer, DPOConfig\n",
    "from arguments import DpoConfig\n",
    "from peft import LoraConfig, TaskType, PeftModel\n",
    "\n",
    "from qwen.modeling_qwen import QWenLMHeadModel\n",
    "from qwen.tokenization_qwen import QWenTokenizer\n",
    "\n",
    "os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'\n",
    "\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from os.path import dirname, abspath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(split: str, file: str, cache_dir: str = '.cache') -> Dataset:\n",
    "    \"\"\"Load the Anthropic Helpful-Harmless dataset from Hugging Face and convert it to the necessary format.\n",
    "\n",
    "    The dataset is converted to a dictionary with the following structure:\n",
    "    {\n",
    "        'prompt': List[str],\n",
    "        'chosen': List[str],\n",
    "        'rejected': List[str],\n",
    "    }\n",
    "    \"\"\"\n",
    "    dataset = load_dataset('parquet', data_files=file,  split=split, cache_dir=cache_dir)\n",
    "\n",
    "    def split_prompt_and_responses(sample: dict) -> Dict[str, str]:\n",
    "        return {\n",
    "            # add an eos token for signal that end of sentence, using in generate.\n",
    "            \"prompt\": f\"{sample['prompt']}<|im_end|>\",\n",
    "            \"chosen\": f\"{sample['chosen']}<|im_end|>\",\n",
    "            \"rejected\": f\"{sample['rejected']}<|im_end|>\",\n",
    "        }\n",
    "\n",
    "    return dataset.map(split_prompt_and_responses).shuffle(2333)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用Lora进行轻量化\n",
    "peft_config = LoraConfig(\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM,  # text 2 text lora model \n",
    "    inference_mode=False, \n",
    "    r=16, \n",
    "    lora_alpha=16, \n",
    "    lora_dropout=0.1, \n",
    "    bias=\"all\",\n",
    "    )\n",
    "\n",
    "dpo_config = DpoConfig()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = QWenTokenizer.from_pretrained(dpo_config.tokenizer_dir)\n",
    "\n",
    "tokenizer.pad_token_id = tokenizer.im_end_id\n",
    "tokenizer.bos_token_id = tokenizer.im_end_id\n",
    "tokenizer.eos_token_id = tokenizer.im_end_id\n",
    "# step 2. 加载SFT模型\n",
    "# model_train, model_ref = None, None\n",
    "# if os.path.isdir(config.sft_model_file):\n",
    "# 传入文件夹则 from_pretrained\n",
    "model_train = QWenLMHeadModel.from_pretrained(dpo_config.sft_model_file)\n",
    "model_ref = QWenLMHeadModel.from_pretrained(dpo_config.sft_model_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载训练数据集\n",
    "dataset = get_dataset(\"train\", file=dpo_config.dpo_train_file)\n",
    "\n",
    "# 加载评估数据集\n",
    "# eval_dataset = get_dataset(\"train\", file=dpo_config.dpo_eval_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置K折交叉验证\n",
    "kfold = KFold(n_splits=15, shuffle=True, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fold_results = []\n",
    "\n",
    "for fold, (train_idx, val_idx) in enumerate(kfold.split(dataset)):\n",
    "    print(f'*** Fold {fold + 1} / {kfold.n_splits}***')\n",
    "\n",
    "    train_dataset = dataset.select(train_idx)\n",
    "    val_dataset = dataset.select(val_idx)\n",
    "\n",
    "    # 设置训练参数\n",
    "    training_args = DPOConfig(\n",
    "        eval_strategy='steps',\n",
    "        save_strategy='steps',\n",
    "        eval_steps=500,\n",
    "        save_steps=500,\n",
    "        per_device_train_batch_size=dpo_config.per_device_train_batch_size,\n",
    "        num_train_epochs=dpo_config.num_train_epochs,\n",
    "        auto_find_batch_size=True,\n",
    "        remove_unused_columns=False,\n",
    "        gradient_accumulation_steps=dpo_config.gradient_accumulation_steps,\n",
    "        learning_rate=dpo_config.learning_rate,\n",
    "        logging_first_step=True,\n",
    "        logging_steps=dpo_config.logging_steps,\n",
    "        save_steps=dpo_config.save_steps,\n",
    "        output_dir=f\"{dpo_config.output_dir}/fold_{fold + 1}\",  # 每个 fold 单独输出\n",
    "        optim=\"adamw_torch\",  # adafactor 不适配所有模型，推荐使用 adamw_torch\n",
    "        report_to=\"tensorboard\",\n",
    "        log_level='info',\n",
    "        warmup_steps=dpo_config.warmup_steps,\n",
    "        bf16=False,\n",
    "        fp16=dpo_config.fp16,\n",
    "        seed=dpo_config.seed,\n",
    "        logging_dir=f\"{dpo_config.log_dir}/fold_{fold + 1}\",  # 为每个 fold 设置日志\n",
    "        load_best_model_at_end=True\n",
    "    )\n",
    "\n",
    "    # 计算指标\n",
    "    def compute_metrics(eval_pred):\n",
    "        from sklearn.metrics import accuracy_score\n",
    "        import numpy as np\n",
    "        logits, labels = eval_pred\n",
    "        predictions = np.argmax(logits, axis=-1)\n",
    "        return {\"accuracy\": accuracy_score(labels, predictions)}\n",
    "\n",
    "    # 如果使用 PEFT，设置 ref_model=None\n",
    "    ref_model = None if peft_config else model_ref\n",
    "\n",
    "    dpo_trainer = DPOTrainer(\n",
    "        model=model_train,\n",
    "        ref_model=ref_model,  # 仅在不使用 LoRA 时传入参考模型\n",
    "        peft_config=peft_config,  # LoRA 微调时使用\n",
    "        args=training_args,\n",
    "        train_dataset=train_dataset,\n",
    "        eval_dataset=val_dataset,\n",
    "        generate_during_eval=True,\n",
    "        is_encoder_decoder=False,  # 重要！Qwen 是自回归模型\n",
    "        compute_metrics=compute_metrics,\n",
    "        processing_class=tokenizer\n",
    "    )\n",
    "\n",
    "    # 开始训练\n",
    "    dpo_trainer.train()\n",
    "\n",
    "    # 评估结果\n",
    "    eval_result = dpo_trainer.evaluate()\n",
    "    print(f\"Fold {fold + 1} Results: {eval_result}\")\n",
    "\n",
    "    fold_results.append(eval_result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 计算 K 折交叉验证的平均结果\n",
    "avg_accuracy = np.mean([result[\"eval_accuracy\"] for result in fold_results])\n",
    "print(f\"Average Accuracy over {kfold.n_splits} folds: {avg_accuracy:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_log = pd.DataFrame(dpo_trainer.state.log_history)\n",
    "log_dir = './logs'\n",
    "if not os.path.exists(log_dir):\n",
    "    os.mkdir(log_dir)\n",
    "loss_log.to_csv(f\"{log_dir}/dpo_train_log_{time.strftime('%Y%m%d-%H%M')}.csv\")\n",
    "    \n",
    "# 10. 保存模型/lora\n",
    "suffixe = '/lora/' if peft_config is not None else '/dpo'\n",
    "model_save_dir = '/'.join(dpo_config.sft_model_file.split('/')[0: -1]) + suffixe\n",
    "\n",
    "dpo_trainer.save_model(model_save_dir)\n",
    "print('save model or lora adapter to: {}'.format(model_save_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = DPOConfig(\n",
    "    eval_strategy='Steps',\n",
    "    save_steps='steps',\n",
    "    eval_steps=500,\n",
    "    per_device_train_batch_size=dpo_config.per_device_train_batch_size,\n",
    "    num_train_epochs=dpo_config.num_train_epochs,\n",
    "    auto_find_batch_size=True,\n",
    "    remove_unused_columns=False,\n",
    "    gradient_accumulation_steps=dpo_config.gradient_accumulation_steps,\n",
    "    learning_rate=dpo_config.learning_rate,\n",
    "    logging_first_step=True,\n",
    "    logging_steps=dpo_config.logging_steps,\n",
    "    output_dir=f\"{dpo_config.output_dir}/fold_{fold + 1}\",  # 每个 fold 单独输出\n",
    "    optim=\"adamw_torch\",  # adafactor 不适配所有模型，推荐使用 adamw_torch\n",
    "    report_to=\"tensorboard\",\n",
    "    log_level='info',\n",
    "    warmup_steps=dpo_config.warmup_steps,\n",
    "    bf16=False,\n",
    "    fp16=dpo_config.fp16,\n",
    "    seed=dpo_config.seed,\n",
    "    logging_dir=f\"{dpo_config.log_dir}/fold_{fold + 1}\",  # 为每个 fold 设置日志\n",
    "    load_best_model_at_end=True\n",
    ")\n",
    "trainer = DPOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
